package com.spzx.report.utils.embedding;

import com.spzx.report.constant.PineconeConst;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.pinecone.PineconeEmbeddingStore;
import lombok.extern.slf4j.Slf4j;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * Package: com.spzx.report.utils
 *
 * @author xiongkunYang
 * @description: 向量数据批量上传工具类, 通过langchain4j api实现
 * @since 2025/6/10 18:43
 */
@Slf4j
public class PineconeUploadUtil {

    private static EmbeddingStore<TextSegment> embeddingStore;

    /***
     * 通过集合上传，前期数据处理
     */
    public static void uploadListToPinecone(List<Map<String, Object>> map, String pineconeIndex, String pineconeApiKey) {

        // 创建向量存储
        embeddingStore = PineconeEmbeddingStore.builder()
                .apiKey(pineconeApiKey)
                .index(pineconeIndex)
                .nameSpace(PineconeConst.NAME_SPACE)
                .build();

        // 分批
        List<String> ids = new ArrayList<>();
        List<Embedding> embeddings = new ArrayList<>();
        List<TextSegment> textSegments = new ArrayList<>();
        // 记录行号
        int lineNum = 0;
        // 成功上传的数量
        int successCount = 0;
        // 失败的数量
        int failedCount = 0;

        for (Map<String, Object> groupKey : map) {
            String cleanText = groupKey.get("key_word").toString().trim();
            if (cleanText.isEmpty()) {
                log.error("数据为空，跳过行 {}: {}", lineNum, cleanText);
                failedCount++;
                continue;
            }

            // 生成文本向量
            Embedding embedding = EmbeddingUtil.getEmbedding(cleanText);
            if (embedding.vector().length == 0) {
                log.error("向量化失败，跳过行 {}: {}", lineNum, cleanText);
                failedCount++;
                continue;
            }

            // 构造 Pinecone 记录
            ids.add("no_" + lineNum++);
            embeddings.add(embedding);

            // 添加 metadata
            TextSegment textSegment = TextSegment.from(cleanText,
                    Metadata.metadata("column", groupKey.get("key_column").toString().trim())
                            .put("word", cleanText)
            );
            textSegments.add(textSegment);

            // 批量上传，每次上传50条
            if (ids.size() >= 50) {
                UploadBatch batch = new UploadBatch(ids, embeddings, textSegments);
                upsertToPinecone(batch);
                successCount += ids.size();
                ids.clear();
                embeddings.clear();
                textSegments.clear();
            }
        }

        // 处理剩余数据
        if (!ids.isEmpty()) {
            UploadBatch batch = new UploadBatch(ids, embeddings, textSegments);
            upsertToPinecone(batch);
            successCount += ids.size();
        }

        log.info("上传完成，共处理 {} 行", lineNum);
        log.info("成功上传：{}", successCount);
        log.info("失败上传：{}", failedCount);
    }

    // 将数据上传至向量库
    private static void upsertToPinecone(UploadBatch batch) {
        try {
            embeddingStore.addAll(batch.ids(), batch.embeddings(), batch.textSegments());
        } catch (Exception e) {
            log.error("上传至 Pinecone 失败", e);
            // 重试
            log.info("正在重试上传...");
            upsertToPinecone(batch);
        }
    }

    // 自定义类用于封装上传批次数据
    private record UploadBatch(
            List<String> ids,
            List<Embedding> embeddings,
            List<TextSegment> textSegments) {
    }
}
